package com.code4j.connect.taos;

import com.code4j.connect.AbstractJDBCService;
import com.code4j.enums.DataSourceTypeEnum;
import com.code4j.pojo.JdbcDbInfo;
import com.code4j.pojo.JdbcMapJavaInfo;
import com.code4j.pojo.JdbcSourceInfo;
import com.code4j.pojo.JdbcTableInfo;
import com.taosdata.jdbc.DatabaseMetaDataResultSet;
import com.taosdata.jdbc.rs.RestfulDatabaseMetaData;
import org.apache.commons.lang3.StringUtils;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.*;

/**
 * @author lwp
 * @date 2023-03-15
 **/
public class TaosJDBCServiceImpl extends AbstractJDBCService {
    public TaosJDBCServiceImpl(final JdbcSourceInfo jdbcSourceInfo) {
        super(jdbcSourceInfo);
    }

    @Override
    protected DataSourceTypeEnum dataSourceTypeEnum() {
        return DataSourceTypeEnum.TAOS;
    }

    @Override
    protected String catalog(String catalog, String dbName) {
        return dbName;
    }

    @Override
    public List<JdbcDbInfo> getAllJdbcDbInfo() {
        Connection connection = null;
        Statement stmt = null;
        try {
            connection = getConnection();
            if (connection == null) {
                return null;
            }
            List<JdbcDbInfo> dbNames = new ArrayList<>();
            stmt = connection.createStatement();
            final ResultSet resultSet = stmt.executeQuery("show databases");
            while (resultSet.next()) {
                final String dbName = resultSet.getString(1);
                if (StringUtils.isNotBlank(dbName)) {
                    JdbcDbInfo jdbcDbInfo = new JdbcDbInfo(dbName, jdbcSourceInfo);
                    dbNames.add(jdbcDbInfo);
                }
            }
            return dbNames;
        } catch (Exception e) {
            log.error("错误！获取数据库信息失败！【{}】", e.getMessage());
        } finally {
            close(connection, stmt, null);
        }
        return null;
    }

    @Override
    public List<JdbcTableInfo> getJdbcTableInfo(final JdbcDbInfo jdbcDbInfo, final String schema) {
        Connection connection = null;
        ResultSet rs = null;
        DatabaseMetaDataResultSet superTables = null;
        try {
            jdbcSourceInfo.setInitDb(jdbcDbInfo.getDbName());
            connection = getConnection();
            if (connection == null) {
                return null;
            }
            String dbName = jdbcDbInfo.getDbName();
            List<JdbcTableInfo> tableInfos = new ArrayList<>();
            RestfulDatabaseMetaData metaData = (RestfulDatabaseMetaData) connection.getMetaData();
            //目录名称; 数据库名; 表名称; 表类型;
            rs = metaData.getTables(this.catalog(connection.getCatalog(), dbName), Optional.ofNullable(schema).orElse(schemaPattern(dbName)), tableNamePattern(), types());
            superTables = (DatabaseMetaDataResultSet) metaData.getSuperTables(this.catalog(connection.getCatalog(), dbName), Optional.ofNullable(schema).orElse(schemaPattern(dbName)), tableNamePattern());
            Set<String> childTableNameSet = new HashSet<>();
            while (superTables.next()) {
                final String stableName = superTables.getString(4);
                if (StringUtils.isNotBlank(stableName)) {
                    childTableNameSet.add(superTables.getString(3));
                }
            }
            while (rs.next()) {
                final String tableName = rs.getString("TABLE_NAME");
                if (childTableNameSet.contains(tableName)) {
                    continue;
                }
                JdbcTableInfo jdbcTableInfo = new JdbcTableInfo();
                jdbcTableInfo.setDbName(dbName);
                jdbcTableInfo.setTableName(tableName);
                jdbcTableInfo.setTableType(rs.getString("TABLE_TYPE"));
                jdbcTableInfo.setRemark(rs.getString("REMARKS"));
                jdbcTableInfo.setJdbcSourceInfo(jdbcDbInfo.getJdbcSourceInfo());
                tableInfos.add(jdbcTableInfo);
            }
            return tableInfos;
        } catch (Exception e) {
            log.error("错误！获取数据库表信息失败【{}】", e.getMessage());
        } finally {
            close(connection, null, rs);
        }
        return null;
    }

    @Override
    public List<JdbcMapJavaInfo> getTableColumnInfo(final String schema, final String dbName, final String tableName) {
        Connection connection = null;
        ResultSet rs = null;
        try {
            connection = getConnection();
            if (connection == null) {
                return null;
            }
            List<JdbcMapJavaInfo> result = new ArrayList<>();
            RestfulDatabaseMetaData meta = (RestfulDatabaseMetaData) connection.getMetaData();
            rs = meta.getColumns(this.catalog(connection.getCatalog(), dbName), Optional.ofNullable(schema).orElse(schemaPattern(dbName)), tableName.trim(), null);
            Set<String> columnCache = new HashSet<>();
            String pk = super.getTablePrimaryKeys(schema, dbName, tableName, connection);
            log.debug("获取表信息参数：数据库{}主键是{}", dbName, pk);
            while (rs.next()) {
                String columnName = rs.getString("COLUMN_NAME");
                if (columnCache.contains(columnName)) {
                    continue;
                }
                columnCache.add(columnName);
                Integer dataTypeNum = rs.getInt("DATA_TYPE");
                String dataType = rs.getString("TYPE_NAME");
                String columnComment = rs.getString("REMARKS");
                JdbcMapJavaInfo jdbcMapJavaInfo = new JdbcMapJavaInfo(columnName, dataType, columnComment, columnName.equalsIgnoreCase(pk), dataTypeNum, this);
                result.add(jdbcMapJavaInfo);
            }
            return result;
        } catch (Exception e) {
            log.error("错误！获取数据库表信息失败【{}】", e.getMessage());
        } finally {
            close(connection, null, rs);
        }
        return null;
    }
}
